import os
import torch
import numpy as np
import pandas as pd
from src.utils.log import infomsg, errormsg


def load_turbine_data(cachepath: str, 
                      dataname: str, 
                      datapath: str, 
                      logfile: str=None, 
                      save_flag: bool=True
                      ):
    """
    加载数据的主方法，先访问缓存，再去csv中加载

    Args:
        cachepath (str): cache文件所在目录
        dataname (str): data的名字，用于区分
        datapath (str): csv形式的数据所在的目录
        logfile (str): log文件的完整路径
        save_flag (bool): 是否将数据保存到cache文件
    Returns:
        torch.Tensor: 加载好的数据
    """
    data = None

    # 尝试从cachefile导入数据
    try:
        cachefile = os.path.join(cachepath, f"data_{dataname}.pt")
        infomsg(logfile, f'尝试从{cachefile}中导入数据')        
        if os.path.exists(cachefile):
            data = torch.load(cachefile)
    except Exception as e:
        errormsg(logfile, str(e))
    
    # 如果cachefile中导入数据失败，从datapath中导入
    if data is None:
        infomsg(logfile, f'尝试从{datapath}中导入数据')
        data = load_from_csv(datapath)
    
    T, N, C = data.shape
    infomsg(logfile, f"数据导入成功，时间T={T}，节点数量N={N}，特征数C={C}")
    
    if save_flag:
        torch.save(data, cachefile)
    return data
    
    
def load_from_csv(datapath):
    """
    给出文件夹，按文件名从小到大顺序加载文件夹里所有的csv文件并stack成一个tensor

    Args:
        datapath(str): csv文件所在的位置
    
    Returns:
        torch.Tensor: 从csv文件中加载得到的数据
    """
    data = []
    files = os.listdir(datapath)

    # 按照文件名的顺序加载
    files.sort()
    for file_name in files:
        if file_name.endswith(".csv"):
            file_path = os.path.join(datapath, file_name)
            df = pd.read_csv(file_path)
            df = df.fillna(0)
            df = df.to_numpy().astype(np.float32)
            data.append(torch.tensor(df))
    
    # 默认读入的数据是(N，time_len，features)，改为(time_len，N，features)
    data = torch.stack(data, dim=0).permute(1, 0, 2)
    return data.contiguous()
